#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <mysql/mysql.h>

#define PORT 8080
#define MAX_CLIENTS 50
#define BUFFER_SIZE 1024

pthread_t threads[MAX_CLIENTS]; // 线程数组
pthread_t admin;
int next_empty = 0; // 下一个空闲槽位的索引
int client_sockets[MAX_CLIENTS];
char names[MAX_CLIENTS][10];

void print_names(void)
{
    int i;
    printf("███Current online user:");
    for (i = 0; i < MAX_CLIENTS; i++)
    {
        if (strlen(names[i]) > 0)
            printf("%s  ", names[i]);
    }
    printf("███\n");
}

void *handle_client(void *data)
{
    MYSQL *mysql;
    MYSQL_RES *result;
    MYSQL_ROW row;
    char query[BUFFER_SIZE];
    int query_status;

    char buffer[BUFFER_SIZE] = {0};
    char response[BUFFER_SIZE] = {0};

    int valread, i, j;
    int send_flag = 0;
    int client_socket = *(int *)data;

    // 连接MySQL数据库
    mysql = mysql_init(NULL);
    if (mysql == NULL)
    {
        printf("Failed to initialize MySQL connection: %s\n", mysql_error(mysql));
        exit(1);
    }
    if (mysql_real_connect(mysql, "localhost", "root", "root", "chat", 0, NULL, 0) == NULL)
    {
        printf("Failed to connect to MySQL server: %s\n", mysql_error(mysql));
        exit(1);
    }

    // 创建用户信息表
    sprintf(query, "CREATE TABLE IF NOT EXISTS user (name VARCHAR(10) NOT NULL, pwd VARCHAR(10) NOT NULL)");
    if (mysql_query(mysql, query) != 0)
    {
        printf("Failed to create friends table: %s\n", mysql_error(mysql));
        exit(1);
    }

    // 创建好友关系表
    sprintf(query, "CREATE TABLE IF NOT EXISTS friends (id INT NOT NULL AUTO_INCREMENT PRIMARY KEY, user1 VARCHAR(10) NOT NULL, user2 VARCHAR(10) NOT NULL)");
    if (mysql_query(mysql, query) != 0)
    {
        printf("Failed to create friends table: %s\n", mysql_error(mysql));
        exit(1);
    }

    while (1)
    {
        if ((valread = read(client_socket, buffer, BUFFER_SIZE)) == 0)
        {
            // 客户端关闭连接
            close(client_socket);
            for (i = 0; i < MAX_CLIENTS; i++)
            {
                if (client_sockets[i] == client_socket)
                {
                    // 从client_sockets数组中移除对应的套接字
                    client_sockets[i] = 0;
                    printf("client %s exit.\n", names[i]);
                    strcpy(names[i], "");
                    print_names();
                    break;
                }
            }
            mysql_close(mysql); // 释放 mysql 连接资源
            pthread_exit(NULL); // 关闭线程
        }
        else
        {
            buffer[valread] = '\0';
            // 解析客户端请求
            char *token;
            token = strtok(buffer, " ");
            printf(">>>token from socket[%d]:%s\n", client_socket, token);
            if (strcmp(token, "login") == 0)
            {
                // 登录请求
                char username[10], pwd[10];
                strcpy(username, strtok(NULL, " "));
                strcpy(pwd, strtok(NULL, " "));
                for (i = 0; i < MAX_CLIENTS; i++)
                {
                    if (client_sockets[i] == client_socket)
                    {
                        strcpy(names[i], username);
                        break;
                    }
                }
                sprintf(query, "SELECT * FROM user WHERE name = '%s' AND pwd = '%s'", username, pwd);
                query_status = mysql_query(mysql, query);
                if (query_status != 0)
                {
                    printf("Query failed: %s\n", mysql_error(mysql));
                    strcpy(response, "Error occur when query your information");
                }
                else
                {
                    result = mysql_store_result(mysql);
                    if (result == NULL)
                    {
                        printf("Failed to store result: %s\n", mysql_error(mysql));
                        strcpy(response, "Error storing result");
                    }
                    else
                    {
                        if (mysql_num_rows(result) > 0)
                        {
                            strcpy(response, "login successfully!");
                        }
                        else
                        {
                            strcpy(response, "login failedw(ﾟДﾟ)w");
                        }
                    }
                }
                send(client_socket, response, strlen(response), 0);
            }
            else if (strcmp(token, "register") == 0)
            {
                // 注册请求
                char username[10], pwd[10];
                strcpy(username, strtok(NULL, " "));
                strcpy(pwd, strtok(NULL, " "));
                sprintf(query, "INSERT INTO user VALUES ('%s', '%s')", username, pwd);
                query_status = mysql_query(mysql, query);
                if (query_status != 0)
                {
                    printf("Failed to insert into user table: %s\n", mysql_error(mysql));
                    strcpy(response, "Error inserting into user table, register failed");
                }
                else
                {
                    strcpy(response, "register successfully");
                }
                send(client_socket, response, strlen(response), 0);
            }
            else if (strcmp(token, "add_friend") == 0)
            {
                // 添加好友请求
                char user1[10], user2[10];
                strcpy(user1, strtok(NULL, " "));
                strcpy(user2, strtok(NULL, " "));
                // 查询是否已经是好友关系
                sprintf(query, "SELECT * FROM friends WHERE (user1 = '%s' AND user2 = '%s') OR (user1 = '%s' AND user2 = '%s')", user1, user2, user2, user1);
                query_status = mysql_query(mysql, query);
                if (query_status != 0)
                {
                    printf("Failed to query friends table: %s\n", mysql_error(mysql));
                    strcpy(response, "Error querying friends table");
                }
                else
                {
                    result = mysql_store_result(mysql);
                    if (result == NULL)
                    {
                        printf("Failed to store result: %s\n", mysql_error(mysql));
                        strcpy(response, "Error storing result");
                    }
                    else
                    {
                        if (mysql_num_rows(result) > 0)
                        {
                            strcpy(response, "Already friends");
                        }
                        else
                        {
                            // 添加好友关系
                            sprintf(query, "INSERT INTO friends (user1, user2) VALUES ('%s', '%s')", user1, user2);
                            query_status = mysql_query(mysql, query);
                            if (query_status != 0)
                            {
                                printf("Failed to insert into friends table: %s\n", mysql_error(mysql));
                                strcpy(response, "Error inserting into friends table");
                            }
                            else
                            {
                                strcpy(response, "Friend added successfully");
                            }
                        }
                        mysql_free_result(result);
                    }
                }
                send(client_socket, response, strlen(response), 0);
            }
            else if (strcmp(token, "send_message") == 0)
            {
                // 发送消息请求
                char sender[10], receiver[10];
                strcpy(sender, strtok(NULL, " "));
                strcpy(receiver, strtok(NULL, " "));
                char *message;
                message = strtok(NULL, "");
                printf(">>>sender:%s, receiver:%s, message:%s\n", sender, receiver, message);
                // 查询是否是好友关系
                sprintf(query, "SELECT * FROM friends WHERE (user1 = '%s' AND user2 = '%s') OR (user1 = '%s' AND user2 = '%s')", sender, receiver, receiver, sender);
                query_status = mysql_query(mysql, query);
                if (query_status != 0)
                {
                    printf("Failed to query friends table: %s\n", mysql_error(mysql));
                    strcpy(response, "Error querying friends table");
                }
                else
                {
                    result = mysql_store_result(mysql);
                    if (result == NULL)
                    {
                        printf("Failed to store result: %s\n", mysql_error(mysql));
                        strcpy(response, "Error storing result");
                    }
                    else
                    {
                        if (mysql_num_rows(result) > 0)
                        {
                            // 是好友关系，发送消息
                            sprintf(response, "%s: %s", sender, message);
                            for (j = 0; j < MAX_CLIENTS; j++)
                            {
                                if (strcmp(names[j], receiver) == 0)
                                {
                                    send(client_sockets[j], response, strlen(response), 0);
                                    send_flag = 1;
                                    break;
                                }
                            }
                            if (send_flag)
                            {
                                strcpy(response, "Message sent successfully");
                                send_flag = 0;
                            }
                            else
                            {
                                strcpy(response, "Message sent failed, your friend is off-line");
                            }
                        }
                        else
                        {
                            strcpy(response, "Not friends, cannot send message");
                        }
                        mysql_free_result(result);
                    }
                }
                send(client_socket, response, strlen(response), 0);
            }
            else if (strcmp(token, "broadcast") == 0)
            {
                // 广播消息请求
                char sender[10];
                strcpy(sender, strtok(NULL, " "));
                char *message;
                message = strtok(NULL, " ");
                sprintf(response, "broadcast from %s: %s", sender, message);
                for (i = 0; i < next_empty; i++)
                {
                    send(client_sockets[i], response, strlen(response), 0);
                }
            }
            else
            {
                // 无效请求
                strcpy(response, "Invalid request");
                send(client_socket, response, strlen(response), 0);
            }
        }
    }
    close(client_socket);
    pthread_exit(NULL);
    mysql_close(mysql);
}

void *handle_admin(void)
{
    printf("admin cmd is running.\n");
    char cmd[10];
    while (1)
    {
        scanf("%s", cmd);
        getchar();
        if (strcmp(cmd, "name") == 0)
        {
            print_names();
        }
        else if (strcmp(cmd, "kick") == 0)
        {
            char kickname[10];
            int j, kickflag = 0;
            printf("Please input name you want to kick:");
            scanf("%s", kickname);
            for (j = 0; j < MAX_CLIENTS; j++)
            {
                if (strcmp(names[j], kickname) == 0)
                {
                    send(client_sockets[j], "disconnect", strlen("disconnect"), 0);
                    strcpy(names[j], "");
                    close(client_sockets[j]);
                    client_sockets[j] = 0;
                    printf("%s was kicked\n", kickname);
                    kickflag = 1;
                    print_names();
                    break;
                }
            }
            if (kickflag == 0)
            {
                printf("The user you want to kick is off-line.\n");
            }
            else
            {
                kickflag = 0;
            }
        }
        else if(strcmp(cmd, "quit") == 0)
        {
            exit(-1);
        }
        else
        {
            printf("Invalid cmd.\n");
        }
    }
}

int main()
{
    int server_fd, new_socket, client_socket, max_fd, client_count = 0, activity, i, valread, sd, err;
    struct sockaddr_in address;
    int addrlen = sizeof(address);

    // 创建服务器套接字
    if ((server_fd = socket(AF_INET, SOCK_STREAM, 0)) == 0)
    {
        perror("socket failed");
        exit(EXIT_FAILURE);
    }
    int opt = 1;
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)))
    {
        perror("setsockopt");
        exit(EXIT_FAILURE);
    }

    // 设置服务器套接字地址信息
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(PORT);
    memset(&(address.sin_zero), 0, 8);

    // 绑定套接字和地址信息
    if (bind(server_fd, (struct sockaddr *)&address, sizeof(address)) < 0)
    {
        perror("bind failed");
        exit(EXIT_FAILURE);
    }

    // 监听服务器套接字
    if (listen(server_fd, MAX_CLIENTS) < 0)
    {
        perror("listen");
        exit(EXIT_FAILURE);
    }
    printf("Server started, waiting for clients...\n");

    err = pthread_create(&admin, NULL, handle_admin, NULL);
    if (err != 0)
    {
        perror("pthread_create");
    }

    while (1)
    {
        // 接受客户端连接请求
        client_socket = accept(server_fd, (struct sockaddr *)&address, (socklen_t *)&addrlen);
        if (client_socket < 0)
        {
            perror("accept");
            continue;
        }

        // 判断是否达到最大客户端连接数
        if (next_empty >= MAX_CLIENTS)
        {
            printf("Reached maximum number of clients, connection refused\n");
            close(client_socket);
            continue;
        }

        // 将新客户端的套接字加入数组中
        client_sockets[next_empty] = client_socket;

        // 启动新线程处理该客户端连接请求
        err = pthread_create(&threads[next_empty], NULL, handle_client, &client_sockets[next_empty]);
        if (err != 0)
        {
            perror("pthread_create");
        }

        // 更新下一个空闲槽位的索引
        next_empty++;
    }
    return 0;
}
